#include <cuda_runtime.h>
#include "geluKernel.h"

// constants for approximating the normal cdf
constexpr const float A = 0.5f;
constexpr const float B = 0.7978845608028654f;   // sqrt(2.0/M_PI)
constexpr const float C = 0.035677408136300125f; // 0.044715 * sqrt(2.0/M_PI)

template<typename T, unsigned int TPB>  // thread per block, [1, 1024]
__global__ void geluKernel(const T* input, T* output, const int len) {
    const unsigned idx = blockIdx.x * TPB + threadIdx.x;
    if (idx < len) {
        const T in  = __ldg(input + idx);
        const T cdf = A + A * tanhf(in * (C * in * in + B));
        output[idx] = in * cdf;
    }
}

int geluKernelLauncher(cudaStream_t stream, 
                        const float* input, 
                        float* output, int len) {
    try {
        const unsigned int blz = 256;
        const unsigned int glz = (len + blz - 1) / blz;
        geluKernel<float, blz><<<glz, blz, 0, stream>>>(input, output, len);
        return 0;
    } catch (std::exception const& e) {
        return -1;
    }
}